(*
 * Humpty-Dumpty - a simple Markov chain based random text generator.
 *
 * author: inky-kinky
 * contact: comamm@yandex.ru
 *
 * This code is distributed under the terms of GPL v.3 license.
 *)

let approx_n = ref 50000

let in_file_handle name f =
  let in_channel = open_in name in
  f in_channel; close_in in_channel 

let out_file_handle name f =
  let out_channel = open_out name in
  f out_channel; close_out out_channel 

let show_stat g =
  Printf.printf 
    "Graph built. Nodes: %i. Arcs: %i.\n" 
    (Graph.n_nodes g) 
    (Graph.n_arcs g)

let m_add_unique_node =
  let m = Hashtbl.create !approx_n in
  fun g x ->
    try
      Hashtbl.find m x
    with Not_found ->
      let n = Graph.add_unique_node g x in
      Hashtbl.add m x n; n

(*
 * Apply function f to every sentence from in_channel (sentences are separated 
 * by regexp [,;://.!?]+).
 *)
let iter_sentences f in_channel =
  let prepare_line x = 
    Str.global_replace (Str.regexp "[\010\013]+") " " x 
  in
  let rec iter_complete x = match x with
    [] -> "" 
    | h::[] -> h 
    | h::t -> f h; iter_complete t 
  in
  let buf = Buffer.create 100 in
  try
    while true do
      let sentences = 
        Buffer.add_string buf (prepare_line (input_line in_channel));
        Buffer.add_char buf ' '; (* line must end with a space. *)
        Str.split (Str.regexp "[,;://.!?]+") (Buffer.contents buf)
      in
      Buffer.clear buf; 
      Buffer.add_string buf (iter_complete sentences)
    done
  with End_of_file -> f (Buffer.contents buf)

(* Add sentence x to graph g. *)
let to_graph g x =
  let rec process prev_n rest_w = match rest_w with
    [] -> ()
    | h::t -> 
      let next_n = m_add_unique_node g h in
      Graph.add_arc g prev_n next_n;
      process next_n t
  in
  let words = Str.split (Str.regexp "[ '()]+") x in
  match words with
    [] -> () 
    | h::[] -> ()
    | h::t -> process (m_add_unique_node g h) t

(* Build graph g from a text file name or from stdin. *)
let build_graph g name =
  let f = fun x -> to_graph g x in
  match name with
    "" -> iter_sentences f stdin; g
    | _ -> in_file_handle name (fun x -> iter_sentences f x); g

(* 
 * Generate random text of n words from graph g and send it to file name or to 
 * stdout.
 *)
let generate_text g n name =
  let f x = 
    let current = ref (Graph.random_node g) in
    for i=1 to n do
      output_string x !current.Graph.data;
      output_char x ' ';
      try current := Graph.random_neighbor g !current
      with Graph.No_such_node -> current := Graph.random_node g
    done 
  in
  match name with
    "" -> f stdout
    | _ -> out_file_handle name f

let () =
  let usage_msg = "Usage: gen <options> <file>\n" in
  let mode = ref "g" in
  let gen_n = ref 100 in
  let file = ref "" in 
  let options = Arg.align [ 
    ("-mode", Arg.String (fun x -> mode:=x), 
      " a (to build graph) or g (to generate text). Default is g."); 
    ("-approx_n", Arg.Int (fun x -> approx_n:=x), 
      " approximate number of arcs in graph. Default is 50000."); 
    ("-gen_n", Arg.Int (fun x -> gen_n:=x), 
      " number of words in text to generate. Default is 100."); 
  ] in
    Arg.parse options (fun x -> file:=x) usage_msg;
    Random.self_init ();
    match !mode with
      "a" -> 
        let g = build_graph (Graph.create !approx_n) !file in
        let out_channel = open_out "graph.dat" in
        Marshal.to_channel out_channel g []; 
        close_out out_channel;
        show_stat g
      | "g" ->
        let in_channel = open_in "graph.dat" in
        let g = Marshal.from_channel in_channel in
        close_in in_channel;
        generate_text g !gen_n !file
      | _ -> Arg.usage options usage_msg
